from fastapi import APIRouter
import os
from model.aoi import Lot, Wafer, Image, Failure
from sqlalchemy import and_, or_
import sqlalchemy
import json
import math
from sqlalchemy import func
from utils.database import SessionLocal
from utils.post_item import AOILotItem, WaferMapItem, WaferMapsItem, WaferStackItem, AOIWaferNGItem, AOILotNGItem, LotReportItem, DieStackItem, WaferDieStackItem
from utils.model_infer import delete_aoi_db
from decimal import Decimal

router = APIRouter()

root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# get all lot


@router.post('/all-lots')
async def get_all_lots(item: AOILotItem):
    # 从数据库中获取所有的 lot 数据
    page_index = item.page_index
    page_size = item.page_size
    lot_name = item.lot_name
    res = {}

    offset_data = page_size * (page_index-1)

    with SessionLocal() as session:
        if lot_name:
            lot_num = session.query(func.count(Lot.id)).filter(
                Lot.lot_name.like("%" + lot_name + "%")).scalar()
            lot_paginate = session.query(Lot).filter(and_(Lot.lot_name.like("%" + lot_name + "%"), Lot.status ==
                                                          'COMPLETED', Lot.del_flag == '0')).order_by(Lot.id.desc()).offset(offset_data).limit(page_size)
        else:
            lot_num = session.query(func.count(Lot.id)).scalar()
            lot_paginate = session.query(Lot).filter(and_(Lot.status == 'COMPLETED', Lot.del_flag == '0')).order_by(
                Lot.id.desc()).offset(offset_data).limit(page_size)

        lots = lot_paginate

        lot_data = []

        for lot in lots:
            lot_item = {}
            lot_item["lot_name"] = lot.lot_name
            lot_item["lot_code"] = lot.lot_code
            lot_item["wafers"] = []
            lot_code = lot.lot_code

            wafers = session.query(Wafer).filter(
                and_(Wafer.lot_code == lot_code, Wafer.del_flag == '0'))
            for wafer in wafers:
                wafer_item = {}
                wafer_item["wafer_id"] = wafer.wafer_id
                wafer_item["wafer_code"] = wafer.wafer_code
                wafer_code = wafer.wafer_code

                wafer_item["image_num"] = session.query(sqlalchemy.func.count(
                    Image.id)).filter_by(wafer_code=wafer_code).scalar()

                wafer_item["ng_num"] = wafer_item["image_num"] - session.query(sqlalchemy.func.count(Image.id)).filter(
                    and_(Image.wafer_code == wafer_code, Image.detection_info == "", Image.segmentation_info == "")).scalar()

                lot_item["wafers"].append(wafer_item)

            lot_data.append(lot_item)

        res["data"] = lot_data
        res["lot_num"] = lot_num

        return res


@router.get('/failure-type')
async def failure_type():
    all_failure_type = {}
    with SessionLocal() as session:
        failures = session.query(Failure).filter_by(del_flag="0")
        for failure in failures:
            all_failure_type[failure.failure_id] = failure.failure_name
        return {"failure_types": all_failure_type}


# get one wafer map data
@router.post('/wafer-map')
async def wafer_map(item: WaferMapItem):

    wafer_code = item.wafer_code

    with SessionLocal() as session:

        wafer = session.query(Wafer).filter_by(wafer_code=wafer_code).first()

        wafer_data = json.loads(wafer.wafer_data)
        wafer_index = wafer.wafer_id

        wafer_dim = (len(wafer_data), len(wafer_data[0]))

        return {"wafermap": wafer_data, "wafer_dim": wafer_dim, "wafer_index": wafer_index}


# get all wafer map data
@router.post('/wafer-maps')
async def wafer_maps(item: WaferMapsItem):
    lot_code = item.lot_code
    page_index = item.page_index
    page_size = item.page_size
    all_failure = item.failure

    count_data = []

    with SessionLocal() as session:

        wafers = session.query(Wafer).filter_by(lot_code=lot_code)

        res = []

        for wafer in wafers:
            wafer_data = json.loads(wafer.wafer_data)
            wafer_index = wafer.wafer_id
            count_item = {}
            count_item["wafer_index"] = wafer_index
            count_item["count"] = {}
            wafer_dim = (len(wafer_data), len(wafer_data[0]))

            for i in range(wafer_dim[0]):
                for j in range(wafer_dim[1]):
                    failure_value = wafer_data[i][j]

                    if failure_value in [0, 1]:
                        continue
                    else:
                        for failure_item in failure_value:
                            fault_type = failure_item["fault_type"]

                            if str(fault_type) in all_failure:
                                if count_item["count"].get(fault_type):
                                    count_item["count"][fault_type] += 1
                                else:
                                    count_item["count"][fault_type] = 1

            wafer_item = {"wafermap": wafer_data,
                          "wafer_dim": wafer_dim, "wafer_index": wafer_index}

            res.append(wafer_item)
            count_data.append(count_item)

        res_per_page = res[(
            page_index - 1) * page_size: page_index * page_size]

        total_page = math.ceil(len(res) / page_size)
        total_num = len(res)

        return {"wafer": res_per_page, "total_page": total_page, "total_num": total_num, "count_data": count_data}


# 获取不同 lot 下的 wafer 的堆叠数据
@router.post('/stack-wafer-map')
async def stack_wafer_map(item: WaferStackItem):
    wafer_codes = item.wafer_codes
    failure = item.failure

    max_wafer_index = 0

    with SessionLocal() as session:

        wafer = session.query(Wafer).filter_by(
            wafer_code=wafer_codes[0]).first()
        wafer_data = json.loads(wafer.wafer_data)

        # 每张 wafer 的尺寸，同一个 lot 的 wafer 尺寸都是一眼的，所以可以只取一个
        wafer_dim = (len(wafer_data), len(wafer_data[0]))

        stack_wafermap = [[0 for i in range(wafer_dim[1])]
                          for j in range(wafer_dim[0])]

        for wafer_code in wafer_codes:
            wafer = session.query(Wafer).filter_by(
                wafer_code=wafer_code).first()
            wafer_data = json.loads(wafer.wafer_data)

            for i in range(wafer_dim[0]):
                for j in range(wafer_dim[1]):
                    wafer_value = wafer_data[i][j]
                    if wafer_value == 0:
                        stack_wafermap[i][j] = 0
                    elif wafer_value == 1:
                        if stack_wafermap[i][j] == 0:
                            stack_wafermap[i][j] = 1
                    else:
                        for fault_data in wafer_value:
                            fault_type = fault_data["fault_type"]
                            if fault_type == int(failure):
                                if stack_wafermap[i][j] == 0:
                                    stack_wafermap[i][j] = 2
                                else:
                                    stack_wafermap[i][j] += 1
                                break
                        else:
                            # 如果有其它错误类型，填充为绿色
                            stack_wafermap[i][j] = 1
                    if stack_wafermap[i][j] > max_wafer_index:
                        max_wafer_index = stack_wafermap[i][j]

        return {"wafermap": stack_wafermap, "wafer_dim": wafer_dim, "all_wafer_index": len(wafer_codes), "max_wafer_index": max_wafer_index}


# 获取不同 lot 下的 wafer 的堆叠数据
@router.post('/stack-wafer-map-die')
async def stack_wafer_map_die(item: WaferDieStackItem):
    wafer_codes = item.wafer_codes
    failure = item.failure
    row = item.row
    col = item.col
    die_dim = item.die_dim
    max_wafer_index = 0

    stack_wafermapdie = [[0 for i in range(die_dim)] for j in range(die_dim)]

    with SessionLocal() as session:

        stack_wafer_die = []

        for wafer_code in wafer_codes:
            wafer = session.query(Wafer).filter_by(
                wafer_code=wafer_code).first()
            wafer_data = json.loads(wafer.wafer_data)

            wafer_value = wafer_data[row][col]

            if isinstance(wafer_value, int):
                continue

            for fault_data in wafer_value:
                fault_type = fault_data["fault_type"]
                if fault_type == int(failure):
                    stack_wafer_die.append(fault_data)
                    relative_fault_x = fault_data["relative_fault_x"]
                    relative_fault_y = fault_data["relative_fault_y"]
                    for i in range(die_dim):
                        for j in range(die_dim):
                            if (1/die_dim)*i < relative_fault_x < (1/die_dim)*(i+1) and (1/die_dim)*j < relative_fault_y < (1/die_dim)*(j+1):
                                stack_wafermapdie[i][j] += 1
                            if stack_wafermapdie[i][j] > max_wafer_index:
                                max_wafer_index = stack_wafermapdie[i][j]

        return {"stack_wafer_die": stack_wafer_die, "stack_wafermapdie": stack_wafermapdie, "max_wafer_index": max_wafer_index}


# 推理一个wafer下的所有图片，并将ng图片和正确图片分开返回
@router.post('/infer-wafer-ng')
async def infer_wafer_ng(item: AOIWaferNGItem):

    with SessionLocal() as session:
        wafer_code = item.wafer_code
        page_index = item.page_index
        ng_page_index = item.ng_page_index
        page_size = item.page_size
        # 整数
        failures = item.failures

        ng_predict_data = []
        predict_data = []
        all_failure_num = {}
        lot_name = None

        wafer = session.query(Wafer).filter_by(wafer_code=wafer_code).first()

        if wafer:

            lot_code = wafer.lot_code
            lot_name = session.query(Lot).filter_by(
                lot_code=lot_code).first().lot_name

            images = session.query(Image).filter_by(wafer_code=wafer_code)

            for image in images:
                item = {}
                item["wafer_id"] = wafer.wafer_id
                item["file_name"] = image.file_name
                item["file_path"] = image.file_path

                # 所有错误信息，为了计算出正常的图片
                item["detection_info"] = []
                item["segmentation_info"] = []

                if all_failure_num.get(image.failure_id):
                    all_failure_num[image.failure_id] += 1
                else:
                    all_failure_num[image.failure_id] = 1

                # 表示正常图片
                if int(image.failure_id) == 1:
                    predict_data.append(item)
                    continue

                if failures and (int(image.failure_id) not in failures):
                    continue

                if image.detection_info:
                    detection_info = json.loads(image.detection_info)
                    for key, value in detection_info.items():

                        for detection in value:
                            bbox = detection[0]
                            distance = None
                            pad_box = None
                            if str(key) == "3":
                                distance = detection[2]
                                pad_box = detection[3]
                            item["detection_info"].append(
                                [key, pad_box, bbox, distance])

                if image.segmentation_info:
                    segmentation_info = json.loads(image.segmentation_info)
                    for key, value in segmentation_info.items():

                        for segmentation in value:
                            points = segmentation[0]
                            all_xy = []
                            bbox = segmentation[2]
                            area_rate = None
                            if str(key) == "3":
                                area_rate = int(Decimal(segmentation[3]).quantize(
                                    Decimal("0.01"), rounding="ROUND_HALF_UP")*100)

                            # points 指的是一张图片的所有轮廓
                            for point in points:
                                # point 指一个 pad 的轮廓
                                point_xy = []
                                for xy in point:
                                    x = xy[0][0] + bbox[0]
                                    y = xy[0][1] + bbox[1]
                                    point_xy.append([x, y])
                                all_xy.append(point_xy)
                            item["segmentation_info"].append(
                                [key, bbox, all_xy, area_rate])

                ng_predict_data.append(item)

        predict_data_per_page = predict_data[(
            page_index - 1) * page_size: page_index * page_size]

        # 如果没有传输 ng_page 索引，则是检查图片，需要返回所有的错误图片
        if ng_page_index:
            ng_predict_data_per_page = ng_predict_data[(
                ng_page_index - 1) * page_size: ng_page_index * page_size]
        else:
            ng_predict_data_per_page = ng_predict_data
        total_page = math.ceil(len(predict_data) / page_size)
        total_num = len(predict_data)

        ng_total_page = math.ceil(len(ng_predict_data) / page_size)
        ng_total_num = len(ng_predict_data)

        class_label = []
        all_failures = session.query(Failure).filter_by(del_flag="0")
        for failure in all_failures:
            label_item = {}
            label_item["failure_name"] = failure.failure_name
            label_item["failure_id"] = failure.failure_id
            label_item["failure_num"] = all_failure_num.get(
                failure.failure_id, 0)
            class_label.append(label_item)

        class_label.append({"failure_name": "正常", "failure_id": 1,
                           "failure_num": all_failure_num.get(1, 0)})

        return {"predict_data": predict_data_per_page, "ng_predict_data": ng_predict_data_per_page, "class_label": class_label, "total_page": total_page, "total_num": total_num, "ng_total_page": ng_total_page, "ng_total_num": ng_total_num, "lot_name": lot_name}

# # 推理一个lot下的所有图片，并全部一起返回
# @router.post('/infer-all')
# def infer_all(item):
#     lot_code = item.get("lot_code")
#     page_index = item.get("page_index")
#     page_size = item.get("page_size")

#     wafers = session.query(Wafer).filter_by(lot_code=lot_code)

#     predict_data = []
#     all_failure_num = {}

#     for wafer in wafers:

#         wafer_code = wafer.wafer_code

#         images = session.query(Image).filter_by(wafer_code=wafer_code)

#         for image in images:
#             item = {}
#             item["wafer_id"] = wafer.wafer_id
#             item["file_name"] = image.file_name
#             item["file_path"] = image.file_path

#             item["detection_info"] = []
#             item["segmentation_info"] = []

#             if image.detection_info:
#                 detection_info = json.loads(image.detection_info)
#                 for key, value in detection_info.items():
#                     if all_failure_num.get(key):
#                         all_failure_num[key] += 1
#                     else:
#                         all_failure_num[key] = 1
#                     for detection in value:
#                         bbox = detection[0]
#                         item["detection_info"].append([key, bbox])

#             if image.segmentation_info:
#                 segmentation_info = json.loads(image.segmentation_info)
#                 for key, value in segmentation_info.items():
#                     if all_failure_num.get(key):
#                         all_failure_num[key] += 1
#                     else:
#                         all_failure_num[key] = 1
#                     for segmentation in value:
#                         point = segmentation[0]
#                         all_xy = []
#                         bbox = segmentation[2]

#                         for xy in point:
#                             x = xy[0][0] + bbox[0]
#                             y = xy[0][1] + bbox[1]
#                             all_xy.append([x, y])
#                         item["segmentation_info"].append([key, bbox, all_xy])

#             predict_data.append(item)

#     predict_data_per_page = predict_data[(
#         page_index - 1) * page_size: page_index * page_size]

#     total_page = math.ceil(len(predict_data) / page_size)
#     total_num = len(predict_data)

#     class_label = []
#     all_failures = session.query(Failure).filter_by(del_flag="0")
#     for failure in all_failures:
#         label_item = {}
#         label_item["failure_name"] = failure.failure_name
#         label_item["failure_id"] = failure.failure_id
#         label_item["failure_num"] = all_failure_num.get(str(failure.failure_id), 0)
#         class_label.append(label_item)

#     return {"predict_data": predict_data_per_page, "class_label": class_label, "total_page": total_page, "total_num": total_num}


# 推理一个lot下的所有图片，并将ng图片和正确图片分开返回
@router.post('/infer-lot-ng')
async def infer_ng(item: AOILotNGItem):
    lot_code = item.lot_code
    page_index = item.page_index
    ng_page_index = item.ng_page_index
    page_size = item.page_size
    # 整数
    failures = item.failures

    with SessionLocal() as session:

        lot_name = session.query(Lot).filter_by(
            lot_code=lot_code).first().lot_name
        wafers = session.query(Wafer).filter_by(lot_code=lot_code)

        ng_predict_data = []
        predict_data = []
        all_failure_num = {}

        for wafer in wafers:

            wafer_code = wafer.wafer_code

            images = session.query(Image).filter_by(wafer_code=wafer_code)

            for image in images:
                item = {}
                item["wafer_id"] = wafer.wafer_id
                item["file_name"] = image.file_name
                item["file_path"] = image.file_path
                item["detection_info"] = []
                item["segmentation_info"] = []

                # 统计不同图片的数量
                if all_failure_num.get(image.failure_id):
                    all_failure_num[image.failure_id] += 1
                else:
                    all_failure_num[image.failure_id] = 1

                # 表示正常图片
                if int(image.failure_id) == 1:
                    predict_data.append(item)
                    continue

                elif int(image.failure_id) not in failures:
                    continue

                if image.detection_info:
                    detection_info = json.loads(image.detection_info)
                    for key, value in detection_info.items():
                        # key:str 错误类型，数字表示
                        for detection in value:
                            bbox = detection[0]
                            distance = None
                            pad_box = None
                            if str(key) == "3":
                                distance = detection[2]
                                pad_box = detection[3]
                            item["detection_info"].append(
                                [key, pad_box, bbox, distance])

                if image.segmentation_info:
                    segmentation_info = json.loads(image.segmentation_info)
                    for key, value in segmentation_info.items():

                        for segmentation in value:
                            points = segmentation[0]
                            all_xy = []
                            bbox = segmentation[2]
                            area_rate = None
                            if str(key) == "3":
                                area_rate = int(Decimal(segmentation[3]).quantize(
                                    Decimal("0.01"), rounding="ROUND_HALF_UP")*100)

                            for point in points:
                                # point 指一个 pad 的轮廓
                                point_xy = []
                                for xy in point:
                                    x = xy[0][0] + bbox[0]
                                    y = xy[0][1] + bbox[1]
                                    point_xy.append([x, y])
                                all_xy.append(point_xy)
                            item["segmentation_info"].append(
                                [key, bbox, all_xy, area_rate])

                ng_predict_data.append(item)

        predict_data_per_page = predict_data[(
            page_index - 1) * page_size: page_index * page_size]
        ng_predict_data_per_page = ng_predict_data[(
            ng_page_index - 1) * page_size: ng_page_index * page_size]

        total_page = math.ceil(len(predict_data) / page_size)
        total_num = len(predict_data)

        ng_total_page = math.ceil(len(ng_predict_data) / page_size)
        ng_total_num = len(ng_predict_data)

        class_label = []
        all_failures = session.query(Failure).filter_by(del_flag="0")

        for failure in all_failures:
            label_item = {}
            label_item["failure_name"] = failure.failure_name
            label_item["failure_id"] = failure.failure_id
            label_item["failure_num"] = all_failure_num.get(
                failure.failure_id, 0)
            class_label.append(label_item)

        class_label.append({"failure_name": "正常", "failure_id": 1,
                           "failure_num": all_failure_num.get(1, 0)})

        return {"predict_data": predict_data_per_page, "ng_predict_data": ng_predict_data_per_page, "class_label": class_label, "total_page": total_page, "total_num": total_num, "ng_total_page": ng_total_page, "ng_total_num": ng_total_num, "lot_name": lot_name}


@router.post('/lot-report')
async def lot_report(item: LotReportItem):
    lot_code = item.lot_code

    with SessionLocal() as session:

        lot = session.query(Lot).filter_by(lot_code=lot_code).first()

        lotName = lot.lot_name

        wafers = session.query(Wafer).filter_by(lot_code=lot_code)

        lot_tabel_data = []
        wafer_tabel_data = []

        lot_tabel_data.append({"key": "Lot Name", "value": lotName})
        lot_tabel_data.append({"key": "Wafers", "value": wafers.count()})

        lot_bad_dice = 0
        lot_good_dice = 0
        all_head_data = {}

        for wafer in wafers:
            wafer_map = json.loads(wafer.wafer_data)
            wafer_index = wafer.wafer_id
            wafer_dim = (len(wafer_map), len(wafer_map[0]))

            # 每个 wafer 的相关统计数据
            wafer_bad_dice = 0
            wafer_good_dice = 0

            all_bin_dice = {}

            all_failures = session.query(Failure).filter_by(del_flag='0')

            for failure in all_failures:
                key = "bin%s_dice" % failure.failure_id
                all_bin_dice[key] = 0
                all_head_data[key] = failure.failure_name

            for i in range(wafer_dim[0]):
                for j in range(wafer_dim[1]):
                    failure_value = wafer_map[i][j]
                    if failure_value == 1:
                        wafer_good_dice += 1
                    elif failure_value == 0:
                        continue
                    else:
                        wafer_bad_dice += 1
                        for failure_item in failure_value:
                            fault_type = failure_item["fault_type"]

                            # for failure in failures:
                            #     if fault_type == failure.failure_id:
                            #         key = "bin%s_dice" % failure.failure_id
                            #         all_bin_dice[key] += 1
                            #         break
                            key = "bin%s_dice" % fault_type
                            all_bin_dice[key] += 1

            lot_bad_dice += wafer_bad_dice
            lot_good_dice += wafer_good_dice

            wafer_total_dice = wafer_bad_dice + wafer_good_dice

            wafer_table_item = {"lot": lotName, "wafer_id": wafer_index,
                                "total_dice": wafer_total_dice, "good_dice": wafer_good_dice, "bad_dice": wafer_bad_dice}

            for key, value in all_bin_dice.items():
                wafer_table_item[key] = value

            wafer_tabel_data.append(wafer_table_item)

        lot_total_dice = lot_good_dice + lot_bad_dice
        lot_tabel_data.append({"key": "Total Dice", "value": lot_total_dice})
        lot_tabel_data.append({"key": "Good Dice", "value": lot_good_dice})
        lot_tabel_data.append({"key": "Bad Dice", "value": lot_bad_dice})

        return {"lot_data": lot_tabel_data, "wafer_data": wafer_tabel_data, "all_head_data": all_head_data}


# 获取一个 wafer 下的所有 die 的堆叠数据
@router.post('/stack-wafer-die')
async def stack_wafer_die(item: DieStackItem):
    wafer_code = item.wafer_code
    failures = item.failures

    with SessionLocal() as session:

        wafer = session.query(Wafer).filter_by(wafer_code=wafer_code).first()

        stack_wafer_die = []

        if wafer:
            wafer_data = json.loads(wafer.wafer_data)

            # 每张 wafer 的尺寸，同一个 lot 的 wafer 尺寸都是一眼的，所以可以只取一个
            wafer_dim = (len(wafer_data), len(wafer_data[0]))

            for i in range(wafer_dim[0]):
                for j in range(wafer_dim[1]):
                    wafer_value = wafer_data[i][j]
                    if (wafer_value != 0) and (wafer_value != 1):
                        for fault_data in wafer_value:
                            fault_type = fault_data["fault_type"]
                            if fault_type in failures:
                                stack_wafer_die.append(fault_data)

        return {"stack_wafer_die": stack_wafer_die}


# 清空单个 lot
@router.get('/lot-delete')
async def lot_delete(lot_name):

    with SessionLocal() as session:
        lot = session.query(Lot).filter_by(lot_name=lot_name).first()
        lot_code = lot.lot_code
        wafers = session.query(Wafer).filter_by(lot_code=lot_code)

        for wafer in wafers:
            wafer_code = wafer.wafer_code
            session.query(Image).filter(
                Image.wafer_code == wafer_code).delete()
            session.delete(wafer)
        session.delete(lot)
        # session.query(Lot).filter(Lot.lot_name==lot_name).delete()
        session.commit()

        return "清空数据库成功"

# 清空 aoi 数据库


@router.get('/lot-delete-all')
async def lot_delete_all():
    with SessionLocal() as session:
        session.query(Lot).delete()
        session.query(Wafer).delete()
        session.query(Image).delete()
        session.commit()

        delete_aoi_db()

        return "清空 aoi 数据库成功"
